"""
用于清洗labelme标记的pig_pose数据
主要纠正左右不分的姿势点位
"""

from matplotlib.font_manager import json_dump, json_load
import numpy as np
import os


def classify_points(vector_start, vector_end, points: tuple):
    """
    用于标记向量线段左侧(右手手心朝向线条,中指指向vector_start方向时,拇指所指的方向)的点和右侧的点
    return True 向量左侧 False 向量右侧
    """
    x1, y1 = vector_start
    x2, y2 = vector_end

    (x, y) = points
    # 计算角度
    angle_oa = np.arctan2(y - y1, x - x1) * (180 / np.pi)
    angle_oe = np.arctan2(y2 - y1, x2 - x1) * (180 / np.pi)

    angle = (angle_oa - angle_oe) % 360

    # 分类
    if 0 <= angle < 180:
        return False
    else:
        return True


def deal_data(path):
    """
    纠正左右不分的姿势点位
    """
    label_json = json_load(path)
    head_point = []
    center_point = []
    tail_point = []
    asides = []
    label_json["imageData"] = None
    for shape in label_json["shapes"]:
        pose_point = shape["label"]
        group_id = shape.get("group_id", 0)
        if pose_point != "pig_pose":
            continue
        if group_id == 0 or group_id == 1:
            head_point = shape["points"]
        elif group_id == 2 or group_id == 3:
            tail_point = shape["points"]
        elif group_id == 8:
            center_point = shape["points"]
        elif group_id > 3 and group_id < 8:
            asides.append(shape)
    if (
        len(asides) == 0
        or len(center_point) == 0
        or len(head_point) == 0
        or len(tail_point) == 0
    ):
        json_dump(label_json, path)
        return
    [cpx, cpy] = center_point[0]
    [hpx, hpy] = head_point[0]
    [tpx, tpy] = tail_point[0]
    for aside in asides:
        [spx, spy] = aside["points"][0]
        # 猪头到中间
        if classify_points((cpx, cpy), (hpx, hpy), (spx, spy)):
            if aside["group_id"] == 5:
                aside["group_id"] = 4
        else:
            if aside["group_id"] == 4:
                aside["group_id"] = 5
        # 中间到猪尾
        if classify_points((tpx, tpy), (cpx, cpy), (spx, spy)):
            if aside["group_id"] == 7:
                aside["group_id"] = 6
        else:
            if aside["group_id"] == 6:
                aside["group_id"] = 7
    json_dump(label_json, path)


data_root = f"C:/Users/kang_/Desktop/data"
count = 0
for root, dirs, files in os.walk(data_root):
    for file in files:
        if not file.endswith(".json"):
            continue
        count += 1
        deal_data(os.path.join(root, file))

print(str(count) + "\n")
